package com.pragmatic.cryptography;

import org.springframework.util.StringUtils;

import javax.crypto.Cipher;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * RSA工具类
 *
 * @author wangdong
 * @since 2025/4/22 9:52
 */
public class RSAUtil {

  // 算法名称
  public static final String KEY_ALGORITHM = "RSA";
  // 定义签名算法
  private static final String KEY_RSA_SIGNATURE = "MD5withRSA";

  // ==============================================================================
  // ================================RSA:生成密钥对==================================
  // ==============================================================================

  /**
   * 初始化公钥、私钥对
   *
   * @throws Exception 生成过程异常
   */
  public static void initKey() throws Exception {
    // 为RSA算法创建一个KeyPairGenerator对象
    KeyPairGenerator kpg = KeyPairGenerator.getInstance("RSA");
    kpg.initialize(2048);
    // 生成密匙对
    KeyPair keyPair = kpg.generateKeyPair();
    // 得到公钥和私钥
    String publicKeyStr = encryptBase64(keyPair.getPublic().getEncoded());
    String privateKeyStr = encryptBase64(keyPair.getPrivate().getEncoded());
    System.out.println("publicKeyStr===" + publicKeyStr);
    System.out.println("privateKeyStr===" + privateKeyStr);
  }

  // ==============================================================================
  // ===========================base64转换字符串和字节数组=============================
  // ==============================================================================

  /**
   * BASE64 解码 （Str -> byte[]）
   *
   * @param key 需要Base64解码的字符串
   * @return 字节数组
   */
  public static byte[] decryptBase64(String key) {
    return Base64.getDecoder().decode(key);
  }

  /**
   * BASE64 编码 （byte[] -> Str）
   *
   * @param key 需要Base64编码的字节数组
   * @return 字符串
   */
  public static String encryptBase64(byte[] key) {
    return new String(Base64.getEncoder().encode(key));
  }

  // ==============================================================================
  // ================================RSA:对内容加解密================================
  // ==============================================================================

  /**
   * 公钥加密(对方公钥加密)
   *
   * @param dataStr 待加密的明文
   * @param publicKeyStr 对方公钥
   * @return 加密后的密文
   */
  public static String encryptByPublic(String dataStr, String publicKeyStr) {
    try {
      // 获得公钥
      byte[] publicKeyBytes = decryptBase64(publicKeyStr);
      // 将公钥由字符串转为UTF-8格式的字节数组
      X509EncodedKeySpec keySpec = new X509EncodedKeySpec(publicKeyBytes);
      // 取得待加密数据
      byte[] data = dataStr.getBytes(StandardCharsets.UTF_8);
      KeyFactory factory = KeyFactory.getInstance(KEY_ALGORITHM);
      PublicKey publicKey = factory.generatePublic(keySpec);
      // 对数据加密
      Cipher cipher = Cipher.getInstance(factory.getAlgorithm());
      cipher.init(Cipher.ENCRYPT_MODE, publicKey);
      // 返回加密后由Base64编码的加密信息
      return encryptBase64(cipher.doFinal(data));
    } catch (Exception e) {
      throw new RuntimeException("RSA加密异常", e);
    }
  }

  /**
   * 私钥解密(己方私钥解密)
   *
   * @param encryptedStr 待解密的密文
   * @param privateKeyStr 己方私钥
   * @return 明文
   */
  public static String decryptByPrivate(String encryptedStr, String privateKeyStr) {
    try {
      // 对私钥解密
      byte[] privateKeyBytes = decryptBase64(privateKeyStr);
      // 获得私钥
      PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(privateKeyBytes);
      // 获得待解密数据
      byte[] data = decryptBase64(encryptedStr);
      KeyFactory factory = KeyFactory.getInstance(KEY_ALGORITHM);
      PrivateKey privateKey = factory.generatePrivate(keySpec);
      // 对数据解密
      Cipher cipher = Cipher.getInstance(factory.getAlgorithm());
      cipher.init(Cipher.DECRYPT_MODE, privateKey);
      // 返回UTF-8编码的解密信息
      return new String(cipher.doFinal(data), StandardCharsets.UTF_8);
    } catch (Exception e) {
      throw new RuntimeException("RSA解密异常", e);
    }
  }

  // ==============================================================================
  // =================================RSA:加签和验签=================================
  // ==============================================================================

  /**
   * 用私钥对原文加签
   *
   * @param originStr 待签名内容
   * @param privateKey 己方私钥
   * @return 签名后的密文
   */
  public static String rsaSign(String originStr, String privateKey) {
    if (StringUtils.isEmpty(originStr)) throw new RuntimeException("签名内容不能为空");
    try {
      // 将私钥加密数据字符串转换为字节数组
      byte[] data = originStr.getBytes();
      // 解密由base64编码的私钥
      byte[] bytes = decryptBase64(privateKey);
      // 构造PKCS8EncodedKeySpec对象
      PKCS8EncodedKeySpec pkcs = new PKCS8EncodedKeySpec(bytes);
      // 指定的加密算法
      KeyFactory factory = KeyFactory.getInstance(KEY_ALGORITHM);
      // 取私钥对象
      PrivateKey key = factory.generatePrivate(pkcs);
      // 用私钥对信息生成数字签名
      Signature signature = Signature.getInstance(KEY_RSA_SIGNATURE);
      signature.initSign(key);
      signature.update(data);
      return encryptBase64(signature.sign());
    } catch (Exception e) {
      throw new RuntimeException("RSA签名异常", e);
    }
  }

  /**
   * 校验数字签名 (对方用公钥验签)
   *
   * @param originStr 待验签密文
   * @param publicKey 对方公钥
   * @param sign 签名数据
   * @return 校验成功返回true，失败返回false
   */
  public static boolean rsaVerify(String originStr, String publicKey, String sign) {
    if (StringUtils.isEmpty(originStr)
        || StringUtils.isEmpty(publicKey)
        || StringUtils.isEmpty(sign)) throw new RuntimeException("验签数据不完整,请检查");
    try {
      // 将私钥加密数据字符串转换为字节数组
      byte[] data = originStr.getBytes();
      // 解密由base64编码的公钥
      byte[] bytes = decryptBase64(publicKey);
      // 构造X509EncodedKeySpec对象
      X509EncodedKeySpec keySpec = new X509EncodedKeySpec(bytes);
      // 指定的加密算法
      KeyFactory factory = KeyFactory.getInstance(KEY_ALGORITHM);
      // 取公钥对象
      PublicKey key = factory.generatePublic(keySpec);
      // 用公钥验证数字签名
      Signature signature = Signature.getInstance(KEY_RSA_SIGNATURE);
      signature.initVerify(key);
      signature.update(data);
      return signature.verify(decryptBase64(sign));
    } catch (Exception e) {
      throw new RuntimeException("RSA验签异常", e);
    }
  }
}
